//
// Created by limo on 2025/9/27.
//

#pragma once

#include <iostream>
#include <string>
#include <map>
#include <memory>
#include <chrono>

#include <boost/asio.hpp>
#include <boost/bind/bind.hpp>
#include <ikcp.h>

namespace net {

    namespace asio = boost::asio;
    using boost::asio::ip::udp;

    // KCP 会话 ID 类型
    using ConvIdType = IUINT32;

    // 消息接收回调：void(const std::string&, std::size_t)
    using ReceiveCallback = std::function<void(const std::string&, std::size_t)>;
    // 错误回调：void(const boost::system::error_code&)
    using ErrorCallback = std::function<void(const boost::system::error_code&)>;

    // 外部 KCP 输出回调（转发到客户端的 UDP socket）
    extern "C" {
        int kcpOutput(const char *buf, int len, ikcpcb *kcp, void *user);
    }

    // 简单的时钟函数 (ikcp 需要毫秒时间戳)
    IUINT32 iclock() {
        return static_cast<IUINT32>(std::chrono::duration_cast<std::chrono::milliseconds>(
                std::chrono::steady_clock::now().time_since_epoch()
        ).count());
    }


    // KcpSession 的核心逻辑与 Server 端类似，这里为 Client 专门实现
    class KcpClientSession : public std::enable_shared_from_this<KcpClientSession> {
    public:
        // 构造函数
        KcpClientSession(asio::io_context& ioContext,
                         const udp::endpoint& serverEndpoint,
                         ConvIdType convId,
                         std::function<void(const char*, int)> udpSendFunction)
                : m_convId(convId),
                  m_serverEndpoint(serverEndpoint),
                  m_udpSendFunction(std::move(udpSendFunction)),
                  m_updateTimer(ioContext)
        {
            // 1. 创建 KCP 控制块
            m_kcpControlBlock = ikcp_create(convId, this);
            m_kcpControlBlock->output = kcpOutput;

            // 2. 设置 KCP 模式
            ikcp_nodelay(m_kcpControlBlock, 1, 10, 2, 1);
            ikcp_wndsize(m_kcpControlBlock, 128, 128);

        }

        ~KcpClientSession() {
            if (m_kcpControlBlock) {
                ikcp_release(m_kcpControlBlock);
                m_kcpControlBlock = nullptr;
            }
            m_updateTimer.cancel();
        }

        void start() {
            startKcpUpdateTimer();
        }

        // --- 公有接口 ---
        void input(const char* data, int size) {
            ikcp_input(m_kcpControlBlock, data, size);
            receiveData();
        }

        int send(const std::string& message) {
            return ikcp_send(m_kcpControlBlock, message.c_str(), message.length());
        }

        void setReceiveCallback(ReceiveCallback cb) { m_onReceive = std::move(cb); }

        ConvIdType getConvId() const { return m_convId; }

    private:
        ConvIdType m_convId;
        udp::endpoint m_serverEndpoint;
        ikcpcb* m_kcpControlBlock;

        asio::steady_timer m_updateTimer;

        // 发送函数：将 KCP 封装的数据发送到 UDP socket (不需要 endpoint，因为目标固定)
        std::function<void(const char*, int)> m_udpSendFunction;

        ReceiveCallback m_onReceive;

        // --- KCP 内部逻辑 ---

        void startKcpUpdateTimer() {
            const auto updateInterval = std::chrono::milliseconds(10);
            m_updateTimer.expires_after(updateInterval);
            m_updateTimer.async_wait(
                    boost::bind(&KcpClientSession::handleKcpUpdate, shared_from_this(), asio::placeholders::error)
            );
        }

        void handleKcpUpdate(const boost::system::error_code& error) {
            if (!error) {
                ikcp_update(m_kcpControlBlock, iclock());
                startKcpUpdateTimer();
            } else if (error != asio::error::operation_aborted) {
                std::cerr << "[KCP Client] Timer error: " << error.message() << std::endl;
            }
        }

        void receiveData() {
            int len;
            while ((len = ikcp_peeksize(m_kcpControlBlock)) > 0) {
                std::string buffer(len, '\0');
                int ret = ikcp_recv(m_kcpControlBlock, buffer.data(), len);

                if (ret > 0 && m_onReceive) {
                    m_onReceive(buffer, ret);
                }
            }
        }

        friend int net::kcpOutput(const char *buf, int len, ikcpcb *kcp, void *user);
    };

    // --- KCP 输出回调的实现 ---
    // KCPClientSession::m_udpSendFunction 在这里被调用
    int kcpOutput(const char *buf, int len, ikcpcb *kcp, void *user) {
        if (user) {
            KcpClientSession* session = static_cast<KcpClientSession*>(user);
            session->m_udpSendFunction(buf, len);
        }
        return 0;
    }

    class KcpClient {
    public:
        // 构造函数
        KcpClient(asio::io_context& ioContext, const std::string& serverAddress, short serverPort, ConvIdType convId)
                : m_socket(ioContext, udp::endpoint(udp::v4(), 0)),
                  m_ioContext(ioContext),
                  m_convId(convId)
        {
            udp::resolver resolver(ioContext);

            // 解析服务器地址
            udp::resolver::results_type endpoints = resolver.resolve(udp::v4(), serverAddress, std::to_string(serverPort));
            m_serverEndpoint = *endpoints.begin();

            // 构造 UDP 发送函数 (用于传递给 KcpClientSession)
            auto udpSendFunc = [this](const char* data, int len) {
                this->udpSend(data, len);
            };

            // 创建 KCP 会话
            m_session = std::make_shared<KcpClientSession>(
                    m_ioContext, m_serverEndpoint, convId, std::move(udpSendFunc)
            );
            m_session->start();
            std::cout << "KcpClient initialized (ConvID: " << m_convId << "). Target: "
                      << m_serverEndpoint.address().to_string() << ":" << m_serverEndpoint.port() << std::endl;

            // 启动异步接收循环
            startReceive();
        }

        // --- 公有接口 ---

        // 转发消息发送到 KcpSession
        void sendMessage(const std::string& message) {
            if (m_session) {
                m_session->send(message);
            }
        }

        // 设置消息回调
        void setReceiveCallback(ReceiveCallback cb) {
            if (m_session) m_session->setReceiveCallback(std::move(cb));
        }

        // 设置错误回调
        void setErrorCallback(ErrorCallback cb) {
            m_onError = std::move(cb);
        }

    private:
        // 成员变量
        udp::socket m_socket;
        udp::endpoint m_serverEndpoint;
        std::array<char, 2048> m_recvBuffer;
        asio::io_context& m_ioContext;

        ConvIdType m_convId;
        std::shared_ptr<KcpClientSession> m_session;

        ErrorCallback m_onError;

        // KCP 输出调用的底层 UDP 裸发
        void udpSend(const char* data, int len) {
            m_socket.async_send_to(
                    asio::buffer(data, len),
                    m_serverEndpoint,
                    [this](const boost::system::error_code& error, std::size_t /*bytesTransferred*/) {
                        if (error && m_onError) {
                            // 如果发送失败，触发错误回调
                            m_onError(error);
                        }
                    }
            );
        }

        // UDP 接收循环
        void startReceive() {
            m_socket.async_receive_from(
                    asio::buffer(m_recvBuffer),
                    m_serverEndpoint, // 接收时仍然需要端点
                    boost::bind(&KcpClient::handleReceive, this,
                                asio::placeholders::error, asio::placeholders::bytes_transferred)
            );
        }

        // UDP 接收处理
        void handleReceive(const boost::system::error_code& error, std::size_t bytesTransferred) {
            if (!error && m_session) {
                // 将收到的原始 UDP 数据包喂给 KCP
                m_session->input(m_recvBuffer.data(), bytesTransferred);
            } else if (error != asio::error::operation_aborted) {
                if (m_onError) m_onError(error);
            }

            // 保持监听
            startReceive();
        }
    };
} // net
